package com.huangkailong.learning.datastructure.tree;

import java.util.Arrays;
import java.util.function.BiFunction;

/**
 * 堆
 *
 * @author huangkl
 * @since 1.0.0
 */
public class Heap<T extends Comparable<T>> {
    private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;
    private static final int DEFAULT_CAPACITY = 10;
    private int size;
    private T[] elementData;
    private final BiFunction<T, T, Boolean> compare;

    public Heap(HeapType type) {
        compare = type == HeapType.MAX_HEAP ? this::gt : this::lt;
    }

    public Heap(HeapType type, int capacity){
        this(type);
        elementData = newDataArray(capacity);
    }

    private T[] newDataArray(int size) {
        return (T[]) new Comparable[size];
    }

    public boolean add(T element){
        ensureCapacityInternal(size+1);
        this.elementData[size++] = element;
        swim(size-1);
        return true;
    }

    /**
     * 上浮算法
     *
     * @param index 开始上浮的位置
     */
    private void swim(int index){
        int currentIndex = index;
        while (!isRoot(currentIndex)){
            int parentIndex = parentIndexOf(currentIndex);
            if(!compare.apply(elementData[currentIndex], elementData[parentIndex])){
                // 最大堆：如果子节点值比父结点值小，则不需要进行处理，直接结束
                // 最小堆：如果子节点值比父结点值大，则不需要进行处理，直接结束
                break;
            }
            swap(currentIndex, parentIndex);
            currentIndex = parentIndex;
        }
    }

    /**
     * 下沉算法
     *
     * @param index 开始下沉的位置
     */
    private void sink(int index){
        int currentIndex = index;
        while (currentIndex < size()){
            int leftChildIndex = leftChildIndex(currentIndex);
            int rightChildIndex = rightChildIndex(currentIndex);

            if(leftChildIndex >= size()){
                // 当前结点无左孩子和右孩子可以直接退出
                break;
            }

            // 找出需要和父结点比较的孩子
            // 最大堆：找出最大的孩子结点
            // 最小堆：找出最小的孩子结点
            int needCompareIndex = leftChildIndex;
            if(rightChildIndex < size()){
                if(!compare.apply(elementData[leftChildIndex], elementData[rightChildIndex])){
                    needCompareIndex = rightChildIndex;
                }
            }
            if(compare.apply(elementData[currentIndex], elementData[needCompareIndex])){
                // 最大堆：如果父节点值比子结点值小，则不需要进行处理，直接结束
                // 最小堆：如果父节点值比子结点值小，则不需要进行处理，直接结束
                break;
            }
            swap(currentIndex, needCompareIndex);
            currentIndex = needCompareIndex;
        }
    }

    private int leftChildIndex(int index){
        return (index+1)*2-1;
    }

    private int rightChildIndex(int index){
        return (index+1)*2;
    }

    private boolean lt(T element1, T element2){
        return element1.compareTo(element2) < 0;
    }

    private boolean gt(T element1, T element2){
        return element1.compareTo(element2) > 0;
    }

    private void swap(int i, int j){
        T tmp = elementData[i];
        elementData[i] = elementData[j];
        elementData[j] = tmp;
    }

    private int parentIndexOf(int i){
        if(isRoot(i)){
            throw new IllegalArgumentException("i == 0 时无父结点");
        }
        return (i+1)/2-1;
    }

    private boolean isRoot(int i){
        return  i == 0;
    }

    private void ensureCapacityInternal(int minCapacity){
        ensureExplicitCapacity(calculateCapacity(elementData, minCapacity));
    }

    private void ensureExplicitCapacity(int minCapacity){
        if(elementData == null || minCapacity - elementData.length > 0){
            grow(minCapacity);
        }
    }

    private int calculateCapacity(T[] elementData, int minCapacity){
        if(elementData == null){
            return Math.max(DEFAULT_CAPACITY, minCapacity);
        }
        return minCapacity;
    }

    private void grow(int minCapacity){
        int elementDataLen = elementData == null ? 0 : elementData.length;
        int newCapacity = elementDataLen + elementDataLen >> 2;
        if(newCapacity - minCapacity < 0){
            newCapacity = minCapacity;
        }
        if(newCapacity - MAX_ARRAY_SIZE > 0){
            newCapacity = hugeCapacity(newCapacity);
        }
        if(elementData == null){
            elementData = newDataArray(newCapacity);
            return;
        }
        elementData = Arrays.copyOf(elementData, newCapacity);
    }

    private int hugeCapacity(int minCapacity){
        if (minCapacity < 0) {
            // overflow
            throw new OutOfMemoryError();
        }
        return (minCapacity > MAX_ARRAY_SIZE) ?
                Integer.MAX_VALUE :
                MAX_ARRAY_SIZE;
    }


    public T remove(){
        if(isEmpty()){
            return null;
        }
        return doRemove();
    }

    private T doRemove() {
        T tmp = elementData[0];
        size--;
        if(isEmpty()){
            return tmp;
        }
        elementData[0] = elementData[size()];
        elementData[size()] = null;
        sink(0);
        return tmp;
    }


    public int size(){
        return size;
    }

    public boolean isEmpty(){
        return size == 0;
    }

}
